from __future__ import annotations

from dataclasses import field,dataclass,astuple
from codesecurity.feature.objects import CommonFeatureSet,Ast
from codesecurity.feature.tfidf import TfidfModule
from codesecurity.tasks.malicious_code_detect.meta import SuperParameter_MCD
from codesecurity.tasks.malicious_code_detect.preprocessing import MCD_Features

import torch.utils.data as torchdata
import numpy as np
import random

# class MaliciousFeatureBuilder:
    
#     def __init__(self,sp:SuperParameter_MCD,id_mapping=None) -> None:
#         if id_mapping is None:
#             id_mapping={}
#         self.sp=sp
#         self.id_mapping=id_mapping
    
#     def build_raw_basic_feature(self,path,label):
#         common_feature=CommonFeatureSet.from_file(path,None)
#         tokens=common_feature.tokens
#         ast=common_feature.ast_object
#         return MaliciousBasicRawFeature(tokens,ast,label,path)

#     def build_basic_feature(self,raw_feature:MaliciousBasicRawFeature,lexical_module:TfidfModule,syntactic_module:TfidfModule):
#         lexical_vec=lexical_module.get_tfidf_vec(raw_feature.tokens,self.sp.word_bag_dim)
#         syntactic_vec=syntactic_module.get_tfidf_vec(self.ast2bag(raw_feature.ast),self.sp.word_bag_dim)
        
#         if raw_feature.label not in self.id_mapping:
#             self.id_mapping[raw_feature.label]=len(self.id_mapping)
        
#         return MaliciousBasicFeature(syntactic_vec,lexical_vec,self.id_mapping[raw_feature.label])    
    
#     def ast2bag(self,ast:Ast):
#         edges=ast.edges
#         return [ast.ast_type(edge).id for edge in edges]
    
# @dataclass
# class MaliciousBasicRawFeature:
#     tokens:list[str]=field(default_factory=list)
#     ast:list[Ast]=field(default_factory=list)
#     label:str=''
#     origin_path:str=""

# @dataclass
# class MaliciousBasicFeature:
#     syntactic_vec:list[float]=field(default_factory=list)
#     lexical_vec:list[float]=field(default_factory=list)
#     label:int=0


class MaliciousDataset(torchdata.Dataset):
    def __init__(self,features:MCD_Features):
        self.features=features 
        self.index_mapping=list(range(len(self.features)))      

    def label_all(self,label,mapping_id=0):
        self.features.label_all(label,mapping_id)

    def shuffle(self,seed):
        random.seed(seed)
        self.index_mapping=list(range(len(self.features)))  
        random.shuffle(self.index_mapping)
        
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self,index):
        index=self.index_mapping[index]
        return self.features[index]

